#pragma once

#include <unordered_map>
#include <mutex>

template <typename Key, typename Value>
class ThreadSafeUnorderedMap
{
public:
    // 插入或更新键值对
    void insert_or_assign(const Key& key, const Value& value) {
        std::lock_guard<std::mutex> lock(_mutex);
        _map[key] = value;
    }

    // 插入键值对，如果键已存在则插入失败
    bool insert(const std::pair<Key, Value>& pair) {
        std::lock_guard<std::mutex> lock(_mutex);
        return _map.insert(pair).second;
    }

    // 查找键对应的迭代器
    auto find(const Key& key) {
        std::lock_guard<std::mutex> lock(_mutex);
        return _map.find(key);
    }

    auto find(const Key& key) const {
        std::lock_guard<std::mutex> lock(_mutex);
        return _map.find(key);
    }

    // 删除键值对
    bool erase(const Key& key) {
        std::lock_guard<std::mutex> lock(_mutex);
        return _map.erase(key) > 0;
    }

    // 检查是否包含某个键
    bool contains(const Key& key) const {
        std::lock_guard<std::mutex> lock(_mutex);
        return _map.find(key) != _map.end();
    }

    // 清空所有键值对
    void clear() {
        std::lock_guard<std::mutex> lock(_mutex);
        _map.clear();
    }

    // 获取当前键值对数量
    std::size_t size() const {
        std::lock_guard<std::mutex> lock(_mutex);
        return _map.size();
    }

    // 支持 operator[]
    Value& operator[](const Key& key) {
        std::lock_guard<std::mutex> lock(_mutex);
        return _map[key];
    }

    const Value& operator[](const Key& key) const {
        std::lock_guard<std::mutex> lock(_mutex);
        return _map.at(key);
    }

    // 返回键值对的迭代器
    auto begin() {
        std::lock_guard<std::mutex> lock(_mutex);
        return _map.begin();
    }

    auto end() {
        std::lock_guard<std::mutex> lock(_mutex);
        return _map.end();
    }

    auto begin() const {
        std::lock_guard<std::mutex> lock(_mutex);
        return _map.begin();
    }

    auto end() const {
        std::lock_guard<std::mutex> lock(_mutex);
        return _map.end();
    }

private:
    std::unordered_map<Key, Value> _map;  // 实际存储数据的unordered_map
    mutable std::mutex _mutex;            // 用于保护unordered_map的互斥锁
};
